/**
 * 
 */
package cs534.crf.test;

import java.io.File;
import java.util.ArrayList;

import org.apache.commons.math3.random.MersenneTwister;

import rddl.EvalException;
import rddl.RDDL;
import rddl.RDDL.DOMAIN;
import rddl.RDDL.INSTANCE;
import rddl.RDDL.LCONST;
import rddl.RDDL.PVAR_INST_DEF;
import rddl.RDDL.PVAR_NAME;
import rddl.State;
import rddl.competition.Client;
import rddl.parser.parser;
import cs534.policy.BoltzmannPolicy;
import cs534.policy.Function2;
import cs534.policy.OnlinePolicyGradient;
import cs534.util.Rddl;

/**
 * @author jhostetler
 *
 */
public class PolicyGradientTwoStates
{
	public static class IndicatorFeature implements Function2<Double, State, ArrayList<PVAR_INST_DEF>>
	{
		private final boolean state_;
		private final boolean action_;
		
		public IndicatorFeature( final boolean state, final boolean action )
		{
			state_ = state;
			action_ = action;
		}
		
		@Override
		public Double apply( final State p0, final ArrayList<PVAR_INST_DEF> p1 )
		{
			final boolean s = (Boolean) p0.getPVariableAssign( new PVAR_NAME( "s" ), new ArrayList<LCONST>() );
			final boolean a = (Boolean) p1.get( 0 )._oValue;
			return (s == state_ && a == action_ ? 1.0 : 0.0);
		}
	}
	
	public static class BasePolicy extends BoltzmannPolicy<BasePolicy, State, ArrayList<PVAR_INST_DEF>>
	{
		private final RDDL rddl_;
		private final MersenneTwister rng_;
		
		public BasePolicy( final double[] parameters, final RDDL rddl )
		{
			this( parameters, rddl, new MersenneTwister() );
		}
		
		public BasePolicy( final double[] parameters, final RDDL rddl, final MersenneTwister rng )
		{
			super( parameters, new IndicatorFeature[] { 
					new IndicatorFeature( false, false ), new IndicatorFeature( false, true ),
					new IndicatorFeature( true, false ), new IndicatorFeature( true, true ) } );
			rddl_ = rddl;
			rng_ = rng;
		}

		public BasePolicy update( final double[] parameters )
		{
			return new BasePolicy( parameters, rddl_, rng_ );
		}
		
		private double[] pi( final State s )
		{
//			System.out.println( "pi" );
			final ArrayList<ArrayList<PVAR_INST_DEF>> actions = Rddl.all_actions( s );
			final double[] pi_s = new double[actions.size()];
			double Z = 0.0;
			int ai = 0;
			for( ArrayList<PVAR_INST_DEF> a : actions ) {
//				System.out.println( "Action: " + a.toString() );
				double q = parameters_[0];
				for( int i = 1; i < parameters_.length; ++i ) {
					q += parameters_[i] * features_[i - 1].apply( s, a );
				}
				pi_s[ai] = Math.exp( q );
				Z += pi_s[ai];
				++ai;
			}
			
			for( ai = 0; ai < pi_s.length; ++ai ) {
				pi_s[ai] /= Z;
			}
			
			return pi_s;
		}
		
		@Override
		protected double marginal_pi( final State s, final int feature )
		{
			final ArrayList<ArrayList<PVAR_INST_DEF>> actions = Rddl.all_actions( s );
			final double[] pi_s = pi( s );
			double result = 0.0;
			int ai = 0;
			for( ArrayList<PVAR_INST_DEF> a : actions ) {
				if( feature == 0 ) {
					result += pi_s[ai];
				}
				else {
					result += pi_s[ai] * features_[feature - 1].apply( s, a );
				}
				++ai;
			}
			
			return result;
		}

		@Override
		public ArrayList<PVAR_INST_DEF> getAction( final State s )
		{
			final double[] pi_s = pi( s );
			final double r = rng_.nextDouble();
			int ai = (r < pi_s[0] ? 0 : 1);
			for( ArrayList<PVAR_INST_DEF> a : Rddl.all_actions( s ) ) {
				--ai;
				if( ai < 0 ) {
					return a;
				}
			}
			return null;
		}
	}
	
	
	
	/**
	 * @param args
	 * @throws Exception 
	 */
	public static void main(String[] args) throws Exception
	{
		final RDDL rddl = parser.parse( new File( "./rddl/cs534_two-states.rddl" ) );
		final DOMAIN d = rddl._tmDomainNodes.values().iterator().next();
		final INSTANCE i = rddl._tmInstanceNodes.values().iterator().next();
		
		final double[] base_policy_parameters = new double[] { 0.1, 0.0, 0.0, 0.0, 0.0 };
		final BasePolicy base_policy = new BasePolicy( base_policy_parameters, rddl );
		Client cl = new Client(d, i, null, new OnlinePolicyGradient<BasePolicy>( base_policy, 0.1, 0.9 ), 
							   "localhost", 3000, "two-states", null);
	}

}
